#ifndef DWF4D_H
#define DWF4D_H
#include "Matrix.h"
#include <cstdlib>
#include <string>
#include <cmath>
#include <iostream>
#include <sstream>
#include <stdexcept>
#include <fstream>
#include <actions/ferm/invert/syssolver_linop_cg_array.h>
#include <complex>
#include <sys/time.h>
#include "dwf_4d_invert.h"
using namespace Chroma;

typedef LatticeFermion TF;
typedef multi1d<LatticeFermion> T5;
typedef multi1d<LatticeColorMatrix> U;



void Gram(LatticeFermion &r, multi1d<LatticeFermion> q, int N){
    for(int i=0;i<N;i++){
        r = r   -  ( innerProduct(q[i],r)/innerProduct(q[i],q[i]) )*q[i];
    }
    r = r*( 1.0/sqrt( norm2(r) ) );
}

void Gram(LatticeFermion &r, multi1d<LatticeFermion> q){
    int N = q.size();
    for(int i=0;i<N;i++){
        r = r   -  ( innerProduct(q[i],r)/innerProduct(q[i],q[i]) )*q[i];
    }
    r = r*( 1.0/sqrt( norm2(r) ) );
}

void Gram(multi1d<LatticeFermion> &q){
    int M = q.size();
    for(int j=0;j<M;j++){
        for(int i=0;i<j;i++){
            q[j] = q[j]   -  ( innerProduct(q[i],q[j])/innerProduct(q[i],q[i]) )*q[i];
        }
        q[j] = q[j]*( 1.0/sqrt( norm2(q[j]) ) );
    }
}

void Gram(vector<LatticeFermion> &q){
    int M = q.size();
    for(int j=0;j<M;j++){
        for(int i=0;i<j;i++){
            q[j] = q[j]   -  ( innerProduct(q[i],q[j])/innerProduct(q[i],q[i]) )*q[i];
        }
        q[j] = q[j]*( 1.0/sqrt( norm2(q[j]) ) );
    }
}


void Gram(multi1d<LatticeFermion> &q, int M){
    for(int j=0;j<M;j++){
        for(int i=0;i<j;i++){
            q[j] = q[j]   -  ( innerProduct(q[i],q[j])/innerProduct(q[i],q[i]) )*q[i];
        }
        q[j] = q[j]*( 1.0/sqrt( norm2(q[j]) ) );
    }
}

void Gram(vector<LatticeFermion> &q, int M){
    for(int j=0;j<M;j++){
        for(int i=0;i<j;i++){
            q[j] = q[j]   -  ( innerProduct(q[i],q[j])/innerProduct(q[i],q[i]) )*q[i];
        }
        q[j] = q[j]*( 1.0/sqrt( norm2(q[j]) ) );
    } 
}



/** Print a 4d Lattice Fermion **/
void print_full_vector(LatticeFermion in, int lx, int ly, int lz, int lt){

    multi1d<int> site(4);	
    site[0] = 0;
    site[1] = 0;
    site[2] = 0;
    site[3] = 0;
    int count = 0;  
    for(int i0=0;i0<lx;i0++){
        for(int i1=0;i1<ly;i1++){
            for(int i2=0;i2<lz;i2++){
                for(int i3=0;i3<lt;i3++){
                    site[0] = i0;site[1] = i1;site[2] = i2;site[3] = i3;

                    //cout << count << " : "; 
                    Fermion F = peekSite(in,site);
                    print_fermion(F);
                    count++;	
                }}}}
    QDPIO::cout << endl;
}



template <class T> class bfmDwf4d : public bfmarg {
    public:
        int Ls;
        T mass;
        int countit;
        bfm_qdp<T> dop;

        double M5;
        multi1d<LatticeColorMatrix> gauge;

        Handle< FermBC<TF,U,U> > *fbc;
        Handle<CreateFermState<TF,U,U> > *cfs;
        UnprecDWFermActArray  *S_f;
        Handle< FermState<TF,U,U> > *fs;
        Handle< UnprecLinearOperatorArray<TF,U,U> > *M5ddwf;




        bool isBoss(void ) { return dop.isBoss() ; };

        void end(void);
        void init(T m, T m5, int N5, multi1d<LatticeColorMatrix> Ugauge, double resid);

        T Munprec(const LatticeFermion &chi,
                LatticeFermion &psi, 
                int dag);
        T Munprec(multi1d<LatticeFermion> &chi_5,
                multi1d<LatticeFermion> &psi_5, 
                int dag);

        T Munprec(Fermion_t chi[2],
                Fermion_t psi[2], 
                int dag);

        T Munprec_TW(Fermion_t chi[2],
                Fermion_t psi[2], 
                int dag);

        T Munprec_TW_QMR(Fermion_t chi[2],
                Fermion_t psi[2], 
                int dag);

        int CGNE_MdagM(LatticeFermion &sol,const LatticeFermion &src);
        int CGNE_Mdag (LatticeFermion &sol,const LatticeFermion &src);
        int CGNE_M    (LatticeFermion &sol,const LatticeFermion &src);

};
template <class T> void bfmDwf4d<T>::end(void)
{
    endQuda();
    delete fbc, cfs, S_f, fs, M5ddwf;
    dop.end();
}

template <class T> void bfmDwf4d<T>::init(T m, T m5, int N5, multi1d<LatticeColorMatrix> Ugauge, double resid)
{
    int device = 0;
    gauge = Ugauge;
    M5 = m5;

    Ls  =N5;
    mass=m;

    // Initialise QUDA and QUDADWF here
    initQuda(0);

    multi1d<int> bcs(Nd);
    bcs[0] = bcs[1] = bcs[2] = bcs[3] = 1;


    fbc = new Handle< FermBC<TF,U,U> > (new SimpleFermBC< TF, U, U >(bcs));

    cfs = new Handle<CreateFermState<TF,U,U> > ( new CreateSimpleFermState<TF,U,U>(*fbc));
    S_f = new UnprecDWFermActArray  (*cfs, M5, mass, Ls);
    fs = new Handle< FermState<TF,U,U> > ( (*S_f).createState(gauge) );

    M5ddwf = new Handle< UnprecLinearOperatorArray<TF,U,U> > ((*S_f).unprecLinOp(*fs,mass));



    countit = 0;
    /**********************************
     * Set up BFM object
     **********************************/

    int lx = QDP::Layout::subgridLattSize()[0];
    int ly = QDP::Layout::subgridLattSize()[1];
    int lz = QDP::Layout::subgridLattSize()[2];
    int lt = QDP::Layout::subgridLattSize()[3];

    bfmarg dwfa;
    //Physics parameters
    dwfa.Ls           = Ls;
    dwfa.M5           = m5;
    dwfa.mass         = mass;
    dwfa.precon_5d    = 1;
    dwfa.max_iter     = 10000;
#ifdef BLEUGENE
    dwfa.time_report_iter = 10000; 
#endif
    dwfa.residual     = resid;
    //Geometry
    dwfa.node_latt[0] = lx;
    dwfa.node_latt[1] = ly;
    dwfa.node_latt[2] = lz;
    dwfa.node_latt[3] = lt;

    multi1d<int> procs = QDP::Layout::logicalSize();
    if ( this -> isBoss() ) {
        printf("dwf_4d: %d dim machine\n\t", procs.size());
    }

    QDPIO::cout << "dwf_4d: required residual = " <<  dwfa.residual << endl;


    for(int mu=0;mu<4;mu++){
        if ( procs[mu]>1 ) {
            dwfa.local_comm[mu] = 0;
        } else { 
            dwfa.local_comm[mu] = 1;
        }
    }
    if ( this -> isBoss() ) {
        printf("dwf_4d: Local comm = ");
        for(int mu=0;mu<4;mu++){
            printf("%d ", dwfa.local_comm[mu]);
        }
        printf("\n");
    }

    multi1d<int> ncoor = QDP::Layout::nodeCoord();

    dop.init(dwfa);
    dop.importGauge(Ugauge);

    dwfa.Ls = 1;
    dwfa.precon_5d = 0;
    dwfa.mass = 1.0;

}

template <class T> int bfmDwf4d<T>::CGNE_M(LatticeFermion &psi,const LatticeFermion &src)
{
    LatticeFermion MdagSrc;
    Munprec(src,MdagSrc,DaggerYes);
    return CGNE_MdagM(psi,MdagSrc);
}
template <class T> int bfmDwf4d<T>::CGNE_Mdag(LatticeFermion &psi,const LatticeFermion &src)
{
    LatticeFermion MinvPsi;
    int iter = CGNE_MdagM(MinvPsi,src);
    Munprec(MinvPsi,psi,DaggerNo);
    return iter;
}
template <class T> int bfmDwf4d<T>::CGNE_MdagM(LatticeFermion &psi,const LatticeFermion &src)
{
    LatticeFermion p;
    LatticeFermion mp;
    LatticeFermion mmp;
    LatticeFermion r;

    Munprec(psi,mp,DaggerNo);
    Munprec(mp,mmp,DaggerYes);

    r = src - mmp;
    p = r;

    Real ssq = norm2(src);
    Real rsq =  residual* residual*ssq;
    Real cp  = norm2(r);

    if ( toDouble(cp) <= toDouble(rsq) ) {
        QDPIO::cout <<"CGNE_unprec converged k=0 - nice guess "<<cp << " " << rsq<<endl;
        return 0;
    }  

    Real c,a,d,b;

    for (int k=1;k<=max_iter;k++){

        c=cp;

        d = Munprec(p,mp,DaggerNo);

        // Update solution vector
        a = c/d;
        psi = a*p+psi;

        // MdagMp
        Munprec(mp,mmp,DaggerYes); 

        r = r-a*mmp;
        cp= norm2(r);
        if ( verbose ) 
            QDPIO::cout <<"dwf_4d: CGNE_unprec k="<<k<<" r^2 = "<<cp<<endl;

        // Stopping condition
        if ( toDouble(cp) <= toDouble(rsq) ) { 
            QDPIO::cout <<"dwf_4d: CGNE_unprec converged k="<<k<<" "<<cp << " " << rsq<<endl;
            Munprec(psi,mp,DaggerNo);
            Munprec(mp,mmp,DaggerYes); 
            r=mmp-src;
            Real true_residual = sqrt(norm2(r)/norm2(src));
            QDPIO::cout <<"dwf_4d: CGNE_unprec true residual "<<true_residual<<endl;
            return k;
        }

        // New (conjugate/M-orthogonal) search direction
        b = cp/c;
        p = b*p+r;
    }
    QDPIO::cout <<"dwf_4d: CGNE_unprec not converged"<<endl;
    return -1;
}

template <class T> T bfmDwf4d<T>::Munprec(const LatticeFermion &chi,
        LatticeFermion &psi, 
        int dag)
{


    multi1d<LatticeFermion> chi_5(Ls), psi_5(Ls);
    /**************************************************************
     * 4d -> 5d system
     **************************************************************/
    // Need chi_5[j] = P_[j][0] chi
    chi_5[0]   = chiralProjectMinus(chi);
    chi_5[Ls-1]= chiralProjectPlus(chi);

    // Zero bulk of chi and result psi
    for(int s=1;s<Ls-1;s++) chi_5[s]=zero;
    for(int s=0;s<Ls;s++)   psi_5[s]=zero;


    Munprec(chi_5,psi_5,dag);


    // Need psi = P^{inv}[0][j] psi_5[j]
    psi = chiralProjectMinus(psi_5[0]);
    psi += chiralProjectPlus (psi_5[Ls-1]);

    T nrm =toDouble(norm2(psi));
    return nrm;
}


template <class T> T bfmDwf4d<T>::Munprec(multi1d<LatticeFermion> &chi_5,
        multi1d<LatticeFermion> &psi_5, 
        int dag)
{

    multi1d<LatticeFermion> tmp(Ls);

    if ( dag ) { 
        /**************************************************************
         * PV^{-d}
         **************************************************************/
        dwf_4d_invert::fullInversion( chi_5, tmp, gauge, M5, DaggerYes);

        /**************************************************************
         * D(m)
         **************************************************************/

        (**M5ddwf)(psi_5, tmp, MINUS); // Chroma

    } else { 
        /**************************************************************
         * D(m)
         **************************************************************/
        (**M5ddwf)(tmp, chi_5, PLUS); // Chroma

        /**************************************************************
         * PV^{-1}
         **************************************************************/
        dwf_4d_invert::fullInversion( tmp, psi_5, gauge, M5, DaggerNo);
    }


}

template <class T> T bfmDwf4d<T>::Munprec(Fermion_t chi[2],
        Fermion_t psi[2], 
        int dag)
{

    multi1d<LatticeFermion> psi_5(Ls), chi_5(Ls);
    for ( int i = 0; i < Ls; i++ )
    {
        psi_5[i] = zero;
        chi_5[i] = zero;
    }

    dop.exportFermion(chi_5, chi[0], 0); dop.exportFermion( chi_5, chi[1], 1);

    Munprec(chi_5, psi_5, dag);

    dop.importFermion(psi_5, psi[0],0); dop.importFermion(psi_5, psi[1],1);


    return (T)1.0;
}


template <class T> T bfmDwf4d<T>::Munprec_TW(Fermion_t chi[2],
        Fermion_t psi[2], 
        int dag)
{

    Munprec(chi,psi,dag);
    return (T)1.0;
}


///Takes same amount of time really...
template <class T> T bfmDwf4d<T>::Munprec_TW_QMR(Fermion_t chi[2],
        Fermion_t psi[2], 
        int dag)
{
    Munprec(chi,psi,dag);
    return (T)1.0;
}







/// q -> q Q
template <class T> void times(multi1d<LatticeFermion> &q, Matrix<T> &Q){
    multi1d<LatticeFermion> S( q.size() );
    for(int j=0;j<q.size();j++){
        S[j] = zero;
        for(int k=0;k<q.size();k++){
            S[j] = S[j] + q[k] * toComplex( Q(k,j) ); 
        }
    }
    for(int j=0;j<q.size();j++){q[j] = S[j];}
}



/// q -> q Q
template <class T> void times(multi1d<LatticeFermion> &q, Matrix<T> &Q, int N){
    multi1d<LatticeFermion> S( N );
    for(int j=0;j<N;j++){
        S[j] = zero;
        for(int k=0;k<N;k++){
            S[j] = S[j] + q[k] * toComplex( Q(k,j) ); 
        }
    }
    for(int j=0;j<N;j++){q[j] = S[j];}
}
#endif
